import numpy as np
import numpy.linalg as LA
import math
#
TYPE_ACC = 1
TYPE_GYR = 2
TYPE_MAG = 3


class MatrixUtils:
    @staticmethod
    def cross(x, y):
        if y.shape[0] == 1 and y.shape[1] == 3:
            z = np.mat([x[0, 1] * y[0, 2] - x[0, 2] * y[0, 1],
                        x[0, 2] * y[0, 0] - x[0, 0] * y[0, 2],
                        x[0, 0] * y[0, 1] - x[0, 1] * y[0, 0]])
        else:
            z = np.mat([x[1, 0] * y[2, 0] - x[2, 0] * y[1, 0],
                        x[2, 0] * y[0, 0] - x[0, 0] * y[2, 0],
                        x[0, 0] * y[1, 0] - x[1, 0] * y[0, 0]]).T
        return z

    @staticmethod
    def createRotateMatrix(axis, rad):
        rotate = [[0 for i in range(3)] for j in range(3)]
        sinr = math.sin(rad)
        cosr = math.cos(rad)

        if axis.shape[0] < axis.shape[1]:
            axis = axis.T

        for i in range(3):
            for j in range(3):
                if i == j:
                    rotate[i][j] = cosr + (1 - cosr) * axis[i, 0] * axis[i, 0]
                if (i + 1) % 3 == j:
                    rotate[i][j] = -sinr * axis[(i + 2) % 3, 0] + (1 - cosr) * axis[i, 0] * axis[j, 0]
                if (i + 2) % 3 == j:
                    rotate[i][j] = sinr * axis[(i + 1) % 3, 0] + (1 - cosr) * axis[i, 0] * axis[j, 0]

        return np.mat(rotate)

    @staticmethod
    def getQuadratureComponent(a, q):
        if a.shape[0] == 1:
            t = (a * q.T)[0, 0]
        else:
            t = (a.T * q)[0, 0]

        tmp = LA.norm(q, 'fro')
        return a - q * float(t / tmp / tmp)


class PostureEstimator:
    _gravity = np.mat([0, 0, 9.81])
    _magnetic = np.mat([1, 0, 0])

    _gravityRotateCount = 0
    _magneticRotateCount = 0

    _localToWorld = np.mat([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
    _worldToLocal = np.mat([[1, 0, 0], [0, 1, 0], [0, 0, 1]])

    _lastAccTime = -1
    _lastMagTime = -1
    _lastGyrTime = -1

    _shouldUpdateLocalToWorldMatrix = False

    def inputRawData(self, type, data, timestamp):
        if type == TYPE_MAG:
            self.inputMag(data, timestamp)
            return
        if type == TYPE_ACC:
            self.inputAcc(data, timestamp)
            return
        if type == TYPE_GYR:
            self.inputGyr(data, timestamp)
            return

    def inputAcc(self, data, timestamp):
        if self._lastAccTime < 0:
            self._lastAccTime = timestamp
            return

        k = 1.0 / (1 + self._gravityRotateCount * 0.001)

        self._gravity = self._gravity * k
        self._gravity = self._gravity + data * (1 - k)

        self._gravityRotateCount = 0
        self._shouldUpdateLocalToWorldMatrix = True
        self._lastAccTime = timestamp

    def inputMag(self, data, timestamp):
        if self._lastMagTime < 0:
            self._lastMagTime = timestamp
            return

        k = 1.0 / (1 + self._magneticRotateCount * 0.01)

        self._magnetic = self._magnetic * k
        self._magnetic = self._magnetic + data * (1 - k)

        self._magneticRotateCount = 0
        self._shouldUpdateLocalToWorldMatrix = True
        self._lastMagTime = timestamp

    def inputGyr(self, data, timestamp):
        if self._lastGyrTime < 0:
            self._lastGyrTime = timestamp
            return

        axis = data
        rad = LA.norm(axis, 'fro') * (timestamp - self._lastGyrTime) / 1e9  # timestamp 1e9?
        axis = axis * (1 / LA.norm(axis, 'fro'))
        rotate = (MatrixUtils.createRotateMatrix(axis, -rad)).T

        self._gravity = self._gravity * rotate
        self._magnetic = self._magnetic * rotate

        self._gravityRotateCount += 1
        self._magneticRotateCount += 1
        self._lastGyrTime = timestamp

    def getLocalToWorldMatrix(self):
        if self._shouldUpdateLocalToWorldMatrix:
            gra2 = self._gravity * (-1 / LA.norm(self._gravity, 'fro'))
            mag2 = MatrixUtils.getQuadratureComponent(self._magnetic, gra2)
            mag2 = mag2 * (1 / LA.norm(mag2, 'fro'))
            y = MatrixUtils.cross(gra2, mag2)
            self._localToWorld = np.mat([[mag2.item(0), mag2.item(1), mag2.item(2)], [y.item(0), y.item(1), y.item(2)], [gra2.item(0), gra2.item(1), gra2.item(2)]])
            self._worldToLocal = self._localToWorld.I

        return self._localToWorld

    def getPosturedData(self, data):
        return np.ravel(self.getLocalToWorldMatrix() * (data.T))

    def getRawData(self, data):
        self.getLocalToWorldMatrix()
        return np.ravel(self._worldToLocal * (data.T))

    def getGravityValue(self):
        return LA.norm(self._gravity, 'fro')

    def getMagneticValue(self):
        return LA.norm(self._magnetic, 'fro')






